# websocket_server.py
# 豆包大模型语音识别WebSocket接口

import base64
import json
import asyncio
from fastapi import WebSocket, WebSocketDisconnect, APIRouter
from .doubao_ws_client import DoubaoASRClient


doubao_ws_stt_router = APIRouter()

@doubao_ws_stt_router.websocket("/ws/speech_v3", name="语音识别WebSocket接口_v3")
async def websocket_endpoint_v3(websocket: WebSocket):
    """
    新版本 WebSocket 接口，客户端发送 JSON：
      {
        "type": "heartbeat"
      }
    或
      {
        "type": "audio",
        "frame_type": 0|1|2,          # 0x00 起始帧, 0x01 中间帧, 0x02 结束帧
        "data": "<base64_encoded_audio_chunk>"
      }
    """
    await websocket.accept()
    print("[v3] WebSocket 连接已接受")
    
    buffer = bytearray()
    client = None
    last_sent_text = ""  # 记录上次发送的文本，避免重复发送

    # 定义结果回调函数
    async def send_result_to_client(text: str):
        nonlocal last_sent_text
        try:
            # 避免重复发送相同的文本
            if text != last_sent_text:
                message = json.dumps({"type": "result", "text": text})
                print(f"[v3] 发送给客户端: {message}")
                await websocket.send_text(message)
                last_sent_text = text
            else:
                print(f"[v3] 跳过重复文本: '{text}'")
        except Exception as e:
            print(f"[v3] 发送结果失败: {e}")

    try:
        while True:
            raw_msg = await websocket.receive_text()
            try:
                msg_obj = json.loads(raw_msg)
            except json.JSONDecodeError as e:
                print(f"[v3] JSON 解析失败: {e}，跳过该消息")
                continue

            msg_type = msg_obj.get("type")
            if msg_type == "heartbeat":
                await websocket.send_text(json.dumps({"type": "pong"}))
                continue
            elif msg_type == "audio":
                frame_type = msg_obj.get("frame_type")
                b64_data = msg_obj.get("data", "")
                try:
                    audio_bytes = base64.b64decode(b64_data)
                except Exception as e:
                    print(f"[v3] base64 解码失败: {e}, 跳过此帧")
                    continue
                if frame_type == 0:
                    buffer.clear()
                    last_sent_text = ""  # 重置上次发送的文本
                    client = DoubaoASRClient(on_result_callback=lambda text: asyncio.create_task(send_result_to_client(text)))
                    await client.connect()
                buffer.extend(audio_bytes)
                # 按 200ms 分包，8k 16bit 单声道 = 1600字节/200ms
                while len(buffer) >= 1600 and client is not None:
                    chunk = bytes(buffer[:1600])
                    await client.send_audio(chunk)
                    buffer = buffer[1600:]
                if frame_type == 2 and client is not None:
                    if len(buffer) > 0:
                        await client.send_audio(bytes(buffer), is_last=True)
                        buffer.clear()
                    await client.close()
                    client = None
                    continue
            else:
                print(f"[v3] 未知消息类型: {msg_type}, 跳过处理")
                continue
    except WebSocketDisconnect:
        print("[v3] 客户端 WebSocket 已断开")
        if client is not None:
            await client.close() 